import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import minimize_scalar
from matplotlib.patches import Ellipse

# ========== 模拟数据生成 ==========
np.random.seed(42)
num_steps = 50
dt = 1.0

# 飞机真实轨迹（匀速运动）
real_x = np.linspace(0, 100, num_steps)
real_y = 0.5 * real_x

# 传感器1 (GPS) 的观测噪声
sensor1_noise = np.random.multivariate_normal(
    [0, 0], [[3.0, 1.5], [1.5, 3.0]], num_steps
)
obs1_x = real_x + sensor1_noise[:, 0]
obs1_y = real_y + sensor1_noise[:, 1]

# 传感器2 (雷达) 的观测噪声
sensor2_noise = np.random.multivariate_normal(
    [0, 0], [[2.0, -1.0], [-1.0, 2.0]], num_steps
)
obs2_x = real_x + sensor2_noise[:, 0]
obs2_y = real_y + sensor2_noise[:, 1]


# ========== 协方差交叉融合函数 ==========
def covariance_intersection(mu1, Sigma1, mu2, Sigma2):
    # 定义优化目标：最小化融合后协方差的迹
    def objective(omega):
        Sigma_inv = omega * np.linalg.inv(Sigma1) + (1 - omega) * np.linalg.inv(Sigma2)
        Sigma = np.linalg.inv(Sigma_inv)
        return np.trace(Sigma)

    # 在 [0, 1] 区间寻找最优 omega
    res = minimize_scalar(objective, bounds=(0, 1), method='bounded')
    omega_opt = res.x

    # 计算融合后的均值和协方差
    Sigma_inv = omega_opt * np.linalg.inv(Sigma1) + (1 - omega_opt) * np.linalg.inv(Sigma2)
    Sigma_fused = np.linalg.inv(Sigma_inv)
    mu_fused = Sigma_fused @ (
            omega_opt * np.linalg.inv(Sigma1) @ mu1 +
            (1 - omega_opt) * np.linalg.inv(Sigma2) @ mu2
    )
    return mu_fused, Sigma_fused, omega_opt


# ========== 逐时刻融合 ==========
fused_means = []
fused_covs = []

for t in range(num_steps):
    # 传感器1的当前观测 (假设已知其协方差)
    mu1 = np.array([obs1_x[t], obs1_y[t]])
    Sigma1 = np.array([[3.0, 1.5], [1.5, 3.0]])

    # 传感器2的当前观测
    mu2 = np.array([obs2_x[t], obs2_y[t]])
    Sigma2 = np.array([[2.0, -1.0], [-1.0, 2.0]])

    # 执行协方差交叉融合
    mu_fused, Sigma_fused, omega = covariance_intersection(mu1, Sigma1, mu2, Sigma2)
    fused_means.append(mu_fused)
    fused_covs.append(Sigma_fused)

fused_means = np.array(fused_means)
fused_covs = np.array(fused_covs)


# ========== 可视化结果 ==========
def plot_ellipse(pos, cov, n_sigma=2, color='r', alpha=0.2):
    """绘制协方差椭圆"""
    eigvals, eigvecs = np.linalg.eigh(cov)
    angle = np.degrees(np.arctan2(eigvecs[1, 0], eigvecs[0, 0]))
    width, height = 2 * n_sigma * np.sqrt(eigvals)
    ellipse = Ellipse(pos, width, height, angle=angle,
                      edgecolor=color, facecolor=color, alpha=alpha)
    plt.gca().add_patch(ellipse)


plt.figure(figsize=(12, 6))

# 真实轨迹
plt.plot(real_x, real_y, 'g-', linewidth=3, label='Real Path')

# 传感器观测
plt.scatter(obs1_x, obs1_y, s=20, c='r', marker='x', label='GPS (Sensor 1)')
plt.scatter(obs2_x, obs2_y, s=20, c='b', marker='+', label='Radar (Sensor 2)')

# 融合结果
plt.plot(fused_means[:, 0], fused_means[:, 1], 'k--', linewidth=2, label='CI Fusion')

# 绘制最后一个时刻的协方差椭圆
plot_ellipse(fused_means[-1], fused_covs[-1], color='purple')
plt.scatter(fused_means[-1, 0], fused_means[-1, 1], c='purple', s=100,
            marker='*', label='Fused Uncertainty')

plt.xlabel('X Position (m)')
plt.ylabel('Y Position (m)')
plt.title('Aircraft Tracking via Covariance Intersection Fusion')
plt.legend()
plt.grid(True)
plt.show()
print('a')